#include <opencv2/opencv.hpp>
#include <iostream>
#include <vector>
#include <cmath>

using namespace cv;
using namespace std;

void Imgupdate(int, void*) {
    ;
}

RotatedRect rectToRotatedRect(const Rect& rect, float angle)
{
    return RotatedRect(
        Point2f(rect.x + rect.width * 0.5f, rect.y + rect.height * 0.5f),
        Size2f(rect.width, rect.height),
        angle
    );
}

// 灯条结构体
struct LightBar {
    RotatedRect rect;
    float length;
    float angle;
    float area;
    float solidity; // 轮廓坚实度
    Point2f top;    // 灯条顶部点
    Point2f bottom; // 灯条底部点
};

// 装甲板结构体
struct ArmorPlate {
    vector<LightBar> lightBars;
    RotatedRect rect;
    Point2f center;
    float confidence;
    int id;         // 装甲板ID
};

// 显示多个图像的函数
void display(vector<Mat>& images) {
    if (images.empty() || images[0].empty()) {
        cerr << "Error: No valid images to display." << endl;
        return;
    }

    int times = images.size() > 3 ? images.size() / 3 : 0;
    int width = images.size() > 3 ? 3 : images.size();
    const int w = images[0].cols * width + 10 * width;
    const int h = images[0].rows * (times + 1) + 10 * (times + 1);
    const int cols = images[0].cols;
    const int rows = images[0].rows;

    if (w <= 0 || h <= 0) {
        cerr << "Error: Invalid image dimensions." << endl;
        return;
    }

    Mat result = Mat::zeros(Size(w, h + 10), images[0].type());
    Rect box(0, 0, cols, rows);

    for (size_t i = 0; i < images.size(); ++i) {
        if (images[i].empty()) {
            cerr << "Warning: Image at index " << i << " is empty." << std::endl;
            continue;
        }
        if (images[i].channels() == 1 && images[0].channels() == 3) {
            cvtColor(images[i], images[i], COLOR_GRAY2BGR);
        }
        int level = i / 3;
        int pos = i % 3;
        box.x = cols * pos + 1 * pos;
        box.y = rows * level + (level >= 1 ? 2 : 0);

        if (box.x + box.width > result.cols || box.y + box.height > result.rows) {
            cerr << "Error: ROI exceeds image boundaries." << endl;
            continue;
        }

        images[i].copyTo(result(box));
    }
    namedWindow("res", WINDOW_AUTOSIZE);
    imshow("res", result);
}

// 计算灯条的顶部和底部点
void calculateLightBarEndpoints(LightBar& lightBar) {
    Point2f vertices[4];
    lightBar.rect.points(vertices);
    
    // 计算四个点的y坐标平均值
    float avgY = 0;
    for (int i = 0; i < 4; i++) {
        avgY += vertices[i].y;
    }
    avgY /= 4.0;
    
    // 根据y坐标将点分为顶部和底部
    vector<Point2f> topPoints, bottomPoints;
    for (int i = 0; i < 4; i++) {
        if (vertices[i].y < avgY) {
            topPoints.push_back(vertices[i]);
        } else {
            bottomPoints.push_back(vertices[i]);
        }
    }
    
    // 计算顶部和底部的中心点
    if (topPoints.size() >= 2) {
        lightBar.top.x = (topPoints[0].x + topPoints[1].x) / 2;
        lightBar.top.y = (topPoints[0].y + topPoints[1].y) / 2;
    }
    
    if (bottomPoints.size() >= 2) {
        lightBar.bottom.x = (bottomPoints[0].x + bottomPoints[1].x) / 2;
        lightBar.bottom.y = (bottomPoints[0].y + bottomPoints[1].y) / 2;
    }
}

// 轮廓检测函数 - 寻找灯条
vector<LightBar> findLightBars(Mat mask, int minArea) {
    vector<LightBar> lightBars;
    vector<vector<Point>> contours;
    vector<Vec4i> hierarchy;
    
    // 查找轮廓
    findContours(mask, contours, hierarchy, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE);
    
    for (const auto& contour : contours) {
        double area = contourArea(contour);
        
        // 面积过滤
        if (area < minArea) continue;
        
        // 计算轮廓的最小外接旋转矩形
        RotatedRect lightRect = minAreaRect(contour);
        
        // 获取轮廓的方向角度（相对于水平方向）
        float angle = lightRect.angle;
        if (lightRect.size.width < lightRect.size.height) {
            angle = angle + 90;
        }
       
        // 计算轮廓坚实度（填充率）
     convexArea = contourArea(convexHull(contour));
        float solidity = (float)(area / convexArea);
       
        // 灯条长宽比过滤
        float aspectRatio = max(lightRect.size.width, lightRect.size.height) / 
                           min(lightRect.size.width, lightRect.size.height);
        
        // 光晕过滤条件
        if (aspectRatio < 3.0 || aspectRatio > 15.0 ||  // 长宽比限制
            solidity < 0.6) {                           // 坚实度过滤光晕
            continue;
        }
        
        // 存储灯条信息
        LightBar lightBar;
        lightBar.rect = lightRect;
        lightBar.length = max(lightRect.size.width, lightRect.size.height);
        lightBar.angle = angle;
        lightBar.area = area;
        lightBar.solidity = solidity;
        
        // 计算灯条的顶部和底部点
        calculateLightBarEndpoints(lightBar);
        
        lightBars.push_back(lightBar);
    }
    
    return lightBars;
}

// 匹配灯条对形成装甲板
vector<ArmorPlate> matchArmorPlates(const vector<LightBar>& lightBars) {
    vector<ArmorPlate> armorPlates;
    
    // 如果灯条数量不足，直接返回
    if (lightBars.size() < 2) return armorPlates;
    
    // 遍历所有灯条对
    for (size_t i = 0; i < lightBars.size(); i++) {
        for (size_t j = i + 1; j < lightBars.size(); j++) {
            const LightBar& left = lightBars[i];
            const LightBar& right = lightBars[j];
            
            // 确保左灯条在右灯条左边（基于x坐标）
            const LightBar& leftBar = (left.rect.center.x < right.rect.center.x) ? left : right;
            const LightBar& rightBar = (left.rect.center.x < right.rect.center.x) ? right : left;
            
            // 计算灯条间的距离
            float distance = norm(leftBar.rect.center - rightBar.rect.center);
            
            // 计算高度差比例
            float heightDiffRatio = abs(leftBar.length - rightBar.length) / max(leftBar.length, rightBar.length);
            
            // 计算角度差（使用优化后的角度计算）
            float angleDiff = abs(leftBar.angle - rightBar.angle);
            if (angleDiff > 90) angleDiff = 180 - angleDiff;
            
            // 计算装甲板角度（使用灯条顶部和底部点）
            float armorAngle = atan2(rightBar.top.y - leftBar.top.y, 
                                    rightBar.top.x - leftBar.top.x) * 180 / CV_PI;
            
            // 计算纵横比（装甲板整体形状）
            float aspectRatio = distance / max(leftBar.length, rightBar.length);
            
            // 计算y轴高度差比例
            float yDiffRatio = abs(leftBar.rect.center.y - rightBar.rect.center.y) / distance;
            
            // 计算顶部和底部的y轴差比例
            float topBottomDiffRatio = abs((rightBar.top.y - leftBar.top.y) - 
                                           (rightBar.bottom.y - leftBar.bottom.y)) / distance;
            
            // 计算坚实度差异
            float solidityDiff = abs(leftBar.solidity - rightBar.solidity);
            
            // 装甲板匹配条件（增强光晕过滤）
            if (heightDiffRatio < 0.4 &&         // 高度差更小
                angleDiff < 30 &&               // 角度差调整
                aspectRatio > 1.0 &&            // 装甲板更宽
                aspectRatio < 8.0 &&            // 装甲板更宽
                yDiffRatio < 0.4 &&             // y轴高度差更小
                topBottomDiffRatio < 0.3 &&     // 顶部和底部对齐
                solidityDiff < 0.2) {           // 坚实度差异小
                
                ArmorPlate armor;
                armor.lightBars.push_back(leftBar);
                armor.lightBars.push_back(rightBar);
                
                // 计算装甲板中心
                armor.center = (leftBar.rect.center + rightBar.rect.center) / 2;
                
                // 创建装甲板的旋转矩形
                float width = distance * 1.1;
                float height = max(leftBar.length, rightBar.length) * 0.8;
                float angle = atan2(rightBar.top.y - leftBar.top.y, 
                                   rightBar.top.x - leftBar.top.x) * 180 / CV_PI;
                
                armor.rect = RotatedRect(armor.center, Size2f(width, height), angle);
                
                // 计算置信度（基于匹配度，增强光晕过滤）
                armor.confidence = (2.0 - heightDiffRatio) * 
                                  (1.5 - angleDiff/30.0) * 
                                  (1.2 - yDiffRatio) *
                                  (1.2 - topBottomDiffRatio) *
                                  (1.3 - solidityDiff);
                
                armorPlates.push_back(armor);
            }
        }
    }
    
    // 按置信度排序（置信度高的优先）
    sort(armorPlates.begin(), armorPlates.end(), 
         [](const ArmorPlate& a, const ArmorPlate& b) {
             return a.confidence > b.confidence;
         });
    
    // 为装甲板分配ID
    for (size_t i = 0; i < armorPlates.size(); i++) {
        armorPlates[i].id = i + 1;
    }
    
    return armorPlates;
}

// 优化红色识别的函数
Mat optimizeRedDetection(const Mat& hsvImg, int h_min1, int h_max1, int h_min2, int h_max2, 
                         int s_min, int s_max, int v_min, int v_max) {
    // 创建两个红色范围的掩码
    Scalar lower1(h_min1, s_min, v_min);
    Scalar upper1(h_max1, s_max, v_max);
    Scalar lower2(h_min2, s_min, v_min);
    Scalar upper2(h_max2, s_max, v_max);
    
    Mat mask1, mask2, mask;
    inRange(hsvImg, lower1, upper1, mask1);
    inRange(hsvImg, lower2, upper2, mask2);
    bitwise_or(mask1, mask2, mask);
    
    // 应用自适应阈值处理，提高红色识别精度
    Mat redChannel, greenChannel, blueChannel;
    vector<Mat> bgrChannels;
    split(hsvImg, bgrChannels);
    redChannel = bgrChannels[2];
    
    Mat adaptiveMask;
    adaptiveThreshold(redChannel, adaptiveMask, 255, ADAPTIVE_THRESH_GAUSSIAN_C, 
                     THRESH_BINARY_INV, 11, 2);
    
    // 合并两种方法的结果
    bitwise_and(mask, adaptiveMask, mask);
    
    // 进一步过滤光晕：使用形态学开运算和面积过滤
    Mat filteredMask;
    morphologyEx(mask, filteredMask, MORPH_OPEN, getStructuringElement(MORPH_ELLIPSE, Size(5, 5)));
    
    return filteredMask;
}

// 绘制识别结果
Mat drawDetectionResult(Mat img, const vector<ArmorPlate>& armorPlates) {
    Mat result = img.clone();
    
    // 绘制装甲板
    for (const auto& armor : armorPlates) {
        // 绘制灯条
        for (const auto& light : armor.lightBars) {
            Point2f vertices[4];
            light.rect.points(vertices);
            for (int i = 0; i < 4; i++) {
                line(result, vertices[i], vertices[(i+1)%4], Scalar(0, 255, 0), 2);
            }
            
            // 标记灯条顶部和底部
            circle(result, light.top, 3, Scalar(255, 255, 0), -1);
            circle(result, light.bottom, 3, Scalar(0, 255, 255), -1);
            
            // 显示灯条坚实度
            putText(result, "S:" + to_string(light.solidity).substr(0, 4),
                    light.top, FONT_HERSHEY_SIMPLEX, 0.4, Scalar(255, 0, 0), 1);
        }
        
        // 绘制装甲板
        Point2f armorVertices[4];
        armor.rect.points(armorVertices);
        for (int i = 0; i < 4; i++) {
            line(result, armorVertices[i], armorVertices[(i+1)%4], Scalar(0, 0, 255), 2);
        }
        
        // 标记装甲板中心
        circle(result, armor.center, 5, Scalar(255, 0, 0), -1);
        
        // 显示装甲板ID和置信度
        putText(result, "Armor #" + to_string(armor.id), 
                Point(armor.center.x - 30, armor.center.y - 20),
                FONT_HERSHEY_SIMPLEX, 0.5, Scalar(255, 0, 0), 2);
        
        putText(result, "Conf: " + to_string(armor.confidence).substr(0, 4), 
                Point(armor.center.x - 30, armor.center.y + 10),
                FONT_HERSHEY_SIMPLEX, 0.5, Scalar(255, 0, 0), 2);
        
        // 显示装甲板角度
        putText(result, "Angle: " + to_string((int)armor.rect.angle) + "°", 
                Point(armor.center.x - 30, armor.center.y + 40),
                FONT_HERSHEY_SIMPLEX, 0.5, Scalar(255, 0, 0), 2);
    }
    
    return result;
}

